import shap
import numpy as np

"""
# 二分类使用示例
print("二分类示例:")
explain_model_with_shap(model, X_test, y_test,
                    feature_names=["特征1", "特征2", "特征3", "特征4"],
                    class_names=["负类", "正类"])

# 多分类使用示例
print("多分类示例:")
explain_model_with_shap(model, X_test, y_test,
                    feature_names=["特征1", "特征2", "特征3", "特征4"],
                    class_names=["类别A", "类别B", "类别C"])

"""


def explain_model_with_shap(
    model,
    X_test,
    y_test=None,
    feature_names=None,
    class_names=None,
    sample_idx=0,
):
    """
    多分类版SHAP模型解释 - 适用于多分类数据集
    :param model: 训练好的XGBoost多分类模型
    :param X_test: 测试集特征
    :param y_test: 测试集标签（可选）
    :param feature_names: 特征名称列表（可选）
    :param class_names: 类别名称列表，如["类别A", "类别B", "类别C"]
    :param sample_idx: 要分析的样本索引
    """
    print("\n=== SHAP模型解释 (多分类版) ===")

    # 创建SHAP解释器
    explainer = shap.TreeExplainer(model)

    # 计算SHAP值（使用少量样本）
    X_test_sample = X_test[:10]  # 只使用前10个样本提高速度
    shap_values = explainer.shap_values(X_test_sample)

    # 获取类别数量
    n_classes = (
        len(shap_values) if isinstance(shap_values, list) else shap_values.shape[0]
    )

    # 设置默认值
    if feature_names is None:
        feature_names = [f"特征_{i+1}" for i in range(X_test.shape[1])]

    if class_names is None:
        class_names = [f"类别_{i}" for i in range(n_classes)]

    # 1. 基础概念解释
    print("\n📊 SHAP基础概念 (多分类):")
    base_values = explainer.expected_value

    print(f"• 类别数量: {n_classes}")
    print("• 每个类别都有一个基准值:")
    for i in range(n_classes):
        base_value = (
            base_values[i]
            if isinstance(base_values, (list, np.ndarray))
            else base_values
        )
        base_prob = log_odds_to_probability(base_value)
        print(
            f"  {class_names[i]}: 基准值={base_value:.4f}, 基准概率={base_prob:.4f} ({base_prob*100:.1f}%)"
        )

    print("   - 含义: 每个类别的平均预测概率")
    print("=" * 50)

    # 2. 全局特征重要性（按类别）
    print("\n🔍 全局特征重要性 (按类别):")
    for class_idx in range(n_classes):
        print(f"\n{class_names[class_idx]} 的特征重要性:")
        class_shap_values = (
            shap_values[class_idx]
            if isinstance(shap_values, list)
            else shap_values[class_idx]
        )
        feature_importance = np.abs(class_shap_values).mean(0)

        for i, name in enumerate(feature_names):
            importance = feature_importance[i]
            print(f"  {name}: {importance:.4f}")

    print("\n💡 解释: 数值越大，该特征对该类别预测的影响越大")
    print("=" * 50)

    # 3. 单个样本解释
    print(f"\n🎯 样本 {sample_idx} 的详细解释 (多分类):")

    # 显示真实标签（如果提供）
    if y_test is not None:
        true_label = y_test[sample_idx]
        print(f"• 真实标签: {true_label} ({class_names[true_label]})")
    else:
        print("• 真实标签: 未提供")

    # 计算每个类别的最终预测概率
    print("\n📈 各类别预测概率:")
    class_probs = []

    for class_idx in range(n_classes):
        base_value = (
            base_values[class_idx]
            if isinstance(base_values, (list, np.ndarray))
            else base_values
        )
        class_shap_values = (
            shap_values[class_idx]
            if isinstance(shap_values, list)
            else shap_values[class_idx]
        )
        sample_shap = class_shap_values[sample_idx]

        final_prediction_log_odds = base_value + np.sum(sample_shap)
        final_prediction_prob = log_odds_to_probability(final_prediction_log_odds)
        class_probs.append(final_prediction_prob)

        print(
            f"  {class_names[class_idx]}: {final_prediction_prob:.4f} ({final_prediction_prob*100:.1f}%)"
        )

    # 找到预测类别
    predicted_class = np.argmax(class_probs)
    max_prob = np.max(class_probs)

    print(f"\n📊 模型预测: {class_names[predicted_class]} ({max_prob*100:.1f}% 置信度)")

    # 4. 特征贡献分析（针对预测类别）
    print(f"\n📈 特征贡献分析 (针对预测类别 {class_names[predicted_class]}):")
    predicted_class_shap = (
        shap_values[predicted_class]
        if isinstance(shap_values, list)
        else shap_values[predicted_class]
    )
    sample_shap = predicted_class_shap[sample_idx]

    for i, name in enumerate(feature_names):
        contribution = sample_shap[i]
        print(f"  {name}: {contribution:+.4f}")

    print("\n💡 特征贡献解释:")
    print(f"• 正值: 增加预测为{class_names[predicted_class]}的概率")
    print(f"• 负值: 减少预测为{class_names[predicted_class]}的概率")

    return explainer, shap_values


def log_odds_to_probability(log_odds):
    """
    将对数几率转换为概率
    :param log_odds: 对数几率值
    :return: 概率值 (0-1之间)
    """
    return 1 / (1 + np.exp(-log_odds))
